import math

import numpy as np
from requests import head
import torch
import torch.nn as nn
from torch import dropout, nn, einsum
import torch.nn.functional as F
from einops import rearrange, repeat
from mmcv.cnn import PLUGIN_LAYERS
from mmcv.cnn import kaiming_init

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim = -1)
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        b, n, _, h = *x.shape, self.heads
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)

        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale

        attn = self.attend(dots)

        out = einsum('b h i j, b h j d -> b h i d', attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

@PLUGIN_LAYERS.register_module()
class RoiAttention(nn.Module):
    """Transformer module for multi-level fpn feature aggregation.

    Args:
        dim (int) : output feature dim
        heads (int) : number of heads of attention module
        dim_head (int) : dim of each head in attention module
    """

    _abbr_ = 'gen_attention_block'

    def __init__(self,
                 dim, 
                 heads = 8, 
                 dim_head = 64,
                 dropout=0.):

        super(RoiAttention, self).__init__()

        # hard range means local range for non-local operation
        inner_dim = dim_head *  heads
        # project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim = -1)
        self.to_k = nn.Conv2d(dim, inner_dim, kernel_size=1, bias = False)
        self.to_q = nn.Conv2d(dim, inner_dim, kernel_size=1, bias = False)
        # self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

        # self.to_out = nn.Sequential(
        #     nn.Linear(inner_dim, dim),
        #     nn.Dropout(dropout)
        # ) if project_out else nn.Identity()

        
        self.init_weights()

    def forward(self, x_input, x_query):
        """
        Args:
            x_input (tensor): [num_rois, num_level, dim, rows, cols]
            x_query (tensor): [num_rois, 1, dim, rows, cols]
        """
        num_rois, num_levels, _ ,rows, cols= x_input.shape
        h = self.heads 
        # 使用原始值作为value
        v = rearrange(x_input, "b n (h d) row col -> b h n (d row col)", h=h) # [heads, num_level * num_rois, head_dim, rows, cols]

        # 计算key 与 query
        x_input_reshape = x_input.reshape(-1, *x_input.shape[2:]) # [num_level * num_rois, dim, rows, cols]
        x_query_reshape = x_query.reshape(-1, *x_query.shape[2:]) # [num_rois, dim, rows, cols]
        x_input_compress = F.adaptive_avg_pool2d(x_input_reshape, 1)
        x_query_compress = F.adaptive_avg_pool2d(x_query_reshape, 1)
        key_input = self.to_k(x_input_compress).squeeze(-1).squeeze(-1) # [num_level * num_rois, dim]
        query = self.to_q(x_query_compress).squeeze(-1).squeeze(-1) # [num_rois, dim]

        key_input = key_input.reshape(num_rois, num_levels, -1)  # [num_rois , num_level, dim]
        query = query.reshape(num_rois, 1, -1) # [num_rois, 1, dim]

        k = rearrange(key_input, 'b n (h d) -> b h n d', h = h)  # [heads, num_level * num_rois, head_dim]
        q = rearrange(query, 'b n (h d) ->b h n d', h = h) # [heads, num_rois, head_dim]
        dots = einsum('b h i d,b h j d -> b h i j', q, k) * self.scale # [heads, num_rois, num_level * num_rois]

        attn = self.attend(dots)
        out = einsum('b h i j, b h j d -> b h i d', attn, v)
        out = rearrange(out, 'b h n (d row col) -> b n (h d) row col', row=rows, col=cols) # n = 1
        out = out.squeeze(1) # [num_rois, dim, row, col]
        return out

    def init_weights(self):
        for m in self.modules():
            if hasattr(m, 'kaiming_init') and m.kaiming_init:
                kaiming_init(
                    m,
                    mode='fan_in',
                    nonlinearity='leaky_relu',
                    bias=0,
                    distribution='uniform',
                    a=1)
